"""
Extract frames from videos into YOLO images folder.

Usage:
  conda activate zny311
  python scripts/extract_frames.py --videos data/local_eval/vids --out datasets/yolo/images/train --fps 2

Then create corresponding labels in datasets/yolo/labels/train with the same stem names.
"""
from __future__ import annotations
import argparse
from pathlib import Path
import cv2


def extract_from_dir(vdir: Path, out_dir: Path, fps: float = 2.0):
    out_dir.mkdir(parents=True, exist_ok=True)
    for vp in sorted(vdir.glob("*.mp4")):
        cap = cv2.VideoCapture(str(vp))
        if not cap.isOpened():
            print(f"[WARN] cannot open {vp}")
            continue
        base = vp.stem
        rate = cap.get(cv2.CAP_PROP_FPS) or 30.0
        step = max(int(rate // max(fps, 0.1)), 1)
        idx = 0
        frame_id = 0
        while True:
            ok, frame = cap.read()
            if not ok:
                break
            if idx % step == 0:
                outp = out_dir / f"{base}_{frame_id:06d}.jpg"
                cv2.imwrite(str(outp), frame)
                frame_id += 1
            idx += 1
        cap.release()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--videos", type=str, required=True)
    ap.add_argument("--out", type=str, required=True)
    ap.add_argument("--fps", type=float, default=2.0)
    args = ap.parse_args()
    extract_from_dir(Path(args.videos), Path(args.out), args.fps)


if __name__ == "__main__":
    main()

